#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
 * @file video_decode.py
 * @author SDK_TEAM
 * @brief
 * @version 0.1
 * @date 2022-11-3
 * Copyright:
 * © 2018 北京灵汐科技有限公司 版权所有。
 * 注意：以下内容均为北京灵汐科技有限公司原创，未经本公司允许，不得转载，否则将视为侵权；对于不遵守此声明或者其他违法使用以下内容者，本公司依法保留追究权。
 * © 2018 Lynxi Technologies Co., Ltd. All rights reserved.
 * NOTICE: All information contained here is, and remains the property of Lynxi. This file can not
 * be copied or distributed without the permission of Lynxi Technologies Co., Ltd.
"""
import sys, getopt
from numpy.core.fromnumeric import repeat
import pylynchipsdk as sdk
import numpy as np
import os
import time

sys.path.append("../../../")
import copy
import common.python.common as common
import common.python.bufferpool as bufferpool
from common.python.blockqueue import block_queue
from common.python.callback_data_struct import recv_cb_data

import threading
from collections import deque

g_exit = False
g_mutex = threading.Lock()
total_send = 0
total_recv = 0
g_running_decoders = 0
g_time_start = 0
g_time_end = 0


class decoder_attr:
    def __init__(self) -> None:
        self.image_path = ""
        self.outnum_min = 1
        self.outnum_max = 10
        self.chan_id = 0
        self.device_id = 0
        self.output_fmt = sdk.lyn_pixel_format_t.LYN_PIX_FMT_NV12
        self.scale = sdk.lyn_scale_t.SCALE_NONE
        self.repeat_num = 1


class jpeg_decoder:
    def __init__(self) -> None:
        self.__attr = ""
        self.__ctx = ""
        self.__send_thread = ""
        self.__recv_thread = ""
        self.__demux_hdl = ""
        self.__jdec_hdl = ""
        self.__jdec_attr = sdk.lyn_jdec_attr_t()
        self.__jdec_info = ""
        self.__send_queue = block_queue()
        self.__send_num = 0
        self.__recv_num = 0

    def __send_thread_func(self):
        # 设置上下文环境 创建发送stream
        sdk.lyn_set_current_context(self.__ctx)
        send_stream, ret = sdk.lyn_create_stream()
        common.error_check(ret, "lyn_create_stream")
        repeat = self.__attr.repeat_num

        file_path = self.__attr.image_path

        while repeat > 0:
            np_jpg = np.fromfile(file_path, dtype=np.byte)
            in_buffer_size = np_jpg.itemsize * np_jpg.size
            packet = sdk.lyn_packet_t()
            packet.size = in_buffer_size
            packet.data = sdk.lyn_numpy_to_ptr(np_jpg)
            ret = sdk.lyn_jdec_send_packet_async(send_stream, self.__jdec_hdl, packet)
            common.error_check(ret, "lyn_jdec_send_packet_async")
            ret = sdk.lyn_synchronize_stream(send_stream)
            common.error_check(ret, "lyn_synchronize_stream")

            # 释放packet内存并通知接收结果
            self.__send_num = self.__send_num + 1
            self.__send_queue.put(self.__send_num)
            repeat = repeat - 1
            global g_mutex
            with g_mutex:
                global total_send
                total_send = total_send + 1
        self.__send_queue.put(-1)

        ret = sdk.lyn_synchronize_stream(send_stream)
        common.error_check(ret, "lyn_synchronize_stream")
        ret = sdk.lyn_destroy_stream(send_stream)
        common.error_check(ret, "lyn_destroy_stream")

    def __recv_thread_func(self):
        # 设置上下文环境 创建接收stream
        sdk.lyn_set_current_context(self.__ctx)
        recv_stream, ret = sdk.lyn_create_stream()
        common.error_check(ret, "lyn_create_stream")
        frame_pool = bufferpool.buffer_pool(self.__jdec_info.output.predict_buf_size, 5)

        global g_time_start
        if g_time_start == 0:
            g_time_start = time.time()

        while 1:
            self.__recv_num = self.__send_queue.take()
            if self.__recv_num == -1:
                break
            cb_data = recv_cb_data()
            cb_data.frame.data = frame_pool.pop()
            cb_data.frame.size = self.__jdec_info.output.predict_buf_size
            cb_data.frame_pool = frame_pool
            cb_data.send_num = self.__send_num
            cb_data.recv_num = self.__recv_num
            cb_data.attr = self.__attr

            # 插入接收指令，并添加接收完成回调函数
            ret = sdk.lyn_jdec_recv_frame_async(
                recv_stream, self.__jdec_hdl, cb_data.frame
            )
            common.error_check(ret, "lyn_jdec_recv_frame_async")
            ret = sdk.lyn_stream_add_callback(
                recv_stream, recv_frame_cb, [cb_data, self.__attr.repeat_num]
            )
            common.error_check(ret, "lyn_stream_add_callback")

        ret = sdk.lyn_synchronize_stream(recv_stream)
        common.error_check(ret, "lyn_synchronize_stream")
        ret = sdk.lyn_destroy_stream(recv_stream)
        common.error_check(ret, "lyn_destroy_stream")

    def run(self, attr: decoder_attr):
        # 创建上下文环境
        self.__attr = copy.copy(attr)
        self.__ctx, ret = sdk.lyn_create_context(attr.device_id)
        common.error_check(ret, "create_context")
        ret = sdk.lyn_register_error_handler(common.default_stream_error_handle)
        common.error_check(ret, "lyn_register_error_handler")

        para = sdk.lyn_image_dec_para_t()
        para.img_url = self.__attr.image_path
        para.scale = sdk.lyn_scale_t.SCALE_NONE
        para.output_fmt = sdk.lyn_pixel_format_t.LYN_PIX_FMT_NV12
        para.align = True
        self.__jdec_info, support, ret = sdk.lyn_image_get_info(para)
        common.error_check(ret, "lyn_image_get_info ")

        if not support:
            print("img url:{} 不支持硬解码".format(para.img_url))
            return

        # 打开解码器
        self.__jdec_attr.output_fmt = self.__attr.output_fmt
        self.__jdec_attr.scale = self.__attr.scale
        self.__jdec_hdl, ret = sdk.lyn_jdec_open(self.__jdec_attr)
        common.error_check(ret, "lyn_jdec_open ")

        # 开启发送线程
        self.__send_thread = threading.Thread(target=self.__send_thread_func)
        self.__send_thread.start()

        # 开启接收线程
        self.__recv_thread = threading.Thread(target=self.__recv_thread_func)
        self.__recv_thread.start()

    def close(self):
        if self.__send_thread and self.__send_thread.is_alive():
            self.__send_thread.join()
            print("self.__send_thread close")
        if self.__recv_thread and self.__recv_thread.is_alive():
            self.__recv_thread.join()
            print("self.__recv_thread close")
        if self.__jdec_hdl != "":
            ret = sdk.lyn_jdec_close(self.__jdec_hdl)
            common.error_check(ret, "lyn_jdec_close")
        if self.__ctx != "":
            ret = sdk.lyn_destroy_context(self.__ctx)
            common.error_check(ret, "lyn_destroy_context")


def recv_frame_cb(cb_data: recv_cb_data):
    # 将内存归还给内存池
    para = cb_data[0]
    repeat_count = cb_data[1]
    para.frame_pool.push(para.frame.data)
    global g_mutex
    global total_recv
    with g_mutex:
        total_recv = total_recv + 1

        if para.recv_num == repeat_count:
            global g_running_decoders
            g_running_decoders = g_running_decoders - 1
            if g_running_decoders == 0:
                global g_time_end
                g_time_end = time.time()

        # host_buf_arr = np.ones(para.frame.size, dtype=np.byte)
        # host_buf = sdk.lyn_numpy_to_ptr(host_buf_arr)
        # ret = sdk.lyn_memcpy(host_buf, para.frame.data, para.frame.size, sdk.lyn_memcpy_dir_t.ServerToClient)
        # common.error_check(ret, "lyn_memcpy")
        # output_file = "{}dump_{:d}_{:0>4d}.yuv".format("./", para.attr.chan_id, para.recv_num)
        # with open(output_file, "wb") as f:
        #     f.write(host_buf_arr.tobytes())
        #     f.close()
    return 0


def running_show():
    global g_exit
    while not g_exit:
        time.sleep(1 / 100)
        print("send:", total_send, " / recv", total_recv, "\r", end="", flush=True)


def usage():
    print("usage: python3 jpeg_decode.py")
    print(
        "\t -i    input file, eg: ../../../data/jpeg_dec_send_packet_sample/test.jpeg"
    )
    print("\t -n    output frame num, default 1,10 eg: num_min,num_max")
    print("\t -d    device id, default 0")
    print("\t -c    channel count, default 1")
    print("\t -r    repeat count, default 1")
    print("\t -h    show help")
    print("\t -?    show help")


def main(argv):
    try:
        opts, args = getopt.getopt(argv, "i:n:d:c:r:h?")
    except getopt.GetoptError:
        usage()
        return 0

    attr = decoder_attr()
    channel_count = 1

    for opt, arg in opts:
        if opt in ("-h", "-?"):
            usage()
            return 0
        elif opt == "-i":
            attr.image_path = arg
        elif opt == "-c":
            channel_count = int(arg)
        elif opt == "-n":
            num = arg.split(",")
            attr.outnum_min = int(num[0])
            attr.outnum_max = int(num[1])
        elif opt == "-d":
            attr.device_id = int(arg)
        elif opt == "-r":
            attr.repeat_num = int(arg)
        else:
            usage()
            return 0

    if attr.image_path == "":
        print("input file miss!!!")
        usage()
        return 0

    running_thread = threading.Thread(target=running_show)
    running_thread.start()

    global g_running_decoders
    g_running_decoders = channel_count

    decoders = deque()
    for i in range(channel_count):
        attr.chan_id = i
        decoder = jpeg_decoder()
        decoder.run(attr)
        decoders.append(decoder)
    for it in decoders:
        it.close()

    global g_exit
    g_exit = True
    running_thread.join()
    global g_time_start
    global g_time_end
    frame_rate = "{:.2f}".format(total_recv / (g_time_end - g_time_start))
    print("decode num ", total_recv, "frame_rate ", frame_rate, " fps")


if __name__ == "__main__":
    main(sys.argv[1:])
